package cn.iocoder.yudao.framework.web.servlet;

import cn.hutool.core.net.URLEncodeUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.digest.MD5;
import cn.hutool.extra.servlet.JakartaServletUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.iocoder.yudao.framework.business.basic.dao.repository.ObjectRedisRepository;
import cn.iocoder.yudao.framework.common.pojo.CommonResult;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.framework.web.core.handler.GlobalExceptionHandler;
import jakarta.annotation.Nonnull;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.access.AccessDeniedException;
import org.springframework.web.servlet.HandlerInterceptor;

import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.getLoginUserId;

/**
 * <pre>
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 *  通过 IP + URL 防止恶意刷新页面、DDoS攻击。
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 * </pre>
 *
 * @author 山野羡民（1032694760@qq.com）
 * @since 2024/12/15
 */
@Slf4j
public class IpUrlLimitInterceptor implements HandlerInterceptor {
    private static final String LOCK_IP_KEY = "xianmin:ip_url_limit:ip_";
    private static final String REQUEST_ONCE_KEY = "xianmin:ip_url_limit:once_";
    private static final long LIMIT_SECONDS = 5; // 多少秒内限制访问多少次
    private static final long LIMIT_ONCE = 500; // 多少秒内限制访问多少次
    private static final int IP_LOCK_SECONDS = 120; // IP 锁定多少秒

    @Override
    public boolean preHandle(@Nonnull HttpServletRequest request, @Nonnull HttpServletResponse response, @Nonnull Object object) {
        String clientIP = JakartaServletUtil.getClientIP(request);
        String uriAndQueryString = request.getRequestURI();
        if (StrUtil.isNotBlank(request.getQueryString())) {
            uriAndQueryString += "?" + request.getQueryString();
        }
        log.debug("请求地址：uriAndQueryString={},ip={},userId={}", uriAndQueryString, clientIP, getLoginUserId());
        ObjectRedisRepository objectRedisRepository = SpringUtil.getBean(ObjectRedisRepository.class);
        String ipHex = MD5.create().digestHex(clientIP);
        String uriHex = URLEncodeUtil.encode(request.getRequestURI());
        String lockIpKey = LOCK_IP_KEY + ipHex;
        String requestOnceKey = REQUEST_ONCE_KEY + ipHex + uriHex;
        objectRedisRepository.setIfAbsent(requestOnceKey, 1L, LIMIT_SECONDS);
        if (objectRedisRepository.hasKey(lockIpKey)) {
            log.info("IP 被临时封禁：{}", clientIP);
            objectRedisRepository.expire(lockIpKey, IP_LOCK_SECONDS);
            responseJson(request, response, "检测到恶意访问，已封禁 IP " + IP_LOCK_SECONDS + "秒");
            return false;
        }
        long totalOnce = objectRedisRepository.increment(requestOnceKey, 1L);
        objectRedisRepository.expire(requestOnceKey, LIMIT_SECONDS);
        if (totalOnce >= LIMIT_ONCE) {
            log.info("恶意访问：uri={},ip={}", request.getRequestURI(), clientIP);
            objectRedisRepository.set(lockIpKey, clientIP, IP_LOCK_SECONDS);
            responseJson(request, response, "访问过于频繁");
            return false;
        }
        return true;
    }

    private void responseJson(HttpServletRequest request, HttpServletResponse response, String msg) {
        GlobalExceptionHandler globalExceptionHandler = SpringUtil.getBean(GlobalExceptionHandler.class);
        CommonResult<?> result = globalExceptionHandler.accessDeniedExceptionHandler(request, new AccessDeniedException(msg));
        JakartaServletUtil.write(response, JsonUtils.toJsonString(result), "application/json; charset=utf-8");
    }

}

